import numpy as np
import pandas as pd
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import AdaBoostClassifier
from sklearn.metrics import accuracy_score
import matplotlib.pyplot as plt


def train_adaboost(dataset, num_base):
    """
    创建并训练AdaBoost集成
    """
    # 创建基学习器（不剪枝的决策树）
    base_learner = DecisionTreeClassifier()
    # 创建AdaBoost集成
    adaboost = AdaBoostClassifier(base_learner, n_estimators=num_base)
    adaboost.fit(dataset.values[:, 0:2], dataset.values[:, -1])
    return adaboost


def plot_boundary(dataset, adaboost, num_base, fig_idx):
    plt.subplot(fig_idx)
    density = dataset['density'].values
    sugar_ratio = dataset['sugar_ratio'].values
    labels = dataset['label'].values

    # 绘制样本点
    plt.scatter(density, sugar_ratio, edgecolors='k', c=labels, cmap=plt.cm.Paired, zorder=10)
    # 绘制分类区域
    ## 确定绘图范围
    x_min = density.min() - 0.1
    x_max = density.max() + 0.1
    y_min = sugar_ratio.min() - 0.1
    y_max = sugar_ratio.max() + 0.1
    ## 将x和y张成二维平面
    XX, YY = np.meshgrid(np.arange(x_min, x_max, 0.001), np.arange(y_min, y_max, 0.001))
    ## 对二维平面内的每一个点进行预测
    ZZ = adaboost.predict(np.c_[XX.ravel(), YY.ravel()])
    ZZ = ZZ.reshape(XX.shape)
    ## 将二维平面按类别进行着色
    plt.pcolormesh(XX, YY, ZZ>0, cmap=plt.cm.Paired, shading='auto')

    plt.title('num_classifier=' + str(num_base))


if __name__ == '__main__':
    # 读取数据
    df = pd.read_csv('watermelon_3_3.csv')
    data = df[['density','sugar_ratio','label']]
    
    # 用不同的基学习器数量训练AdaBoost
    for num_base, fig_idx in zip([1,3,5,11], [221,222,223,224]):
        ada = train_adaboost(data, num_base=num_base)
        labels = data.values[:, -1]
        predicts = ada.predict(data.values[:, 0:2])
        print('\nnum_base = ', num_base)
        print('labels: ', labels)
        print('predict:', predicts)
        print('accuracy:', accuracy_score(labels, predicts))
        plot_boundary(data, ada, num_base, fig_idx)

    plt.tight_layout()
    plt.savefig('adaboost.png')
    plt.show()
